package com.qsl.redisson;

import com.google.common.collect.Lists;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.redisson.Redisson;
import org.redisson.api.RScript;
import org.redisson.api.RedissonClient;
import org.redisson.client.codec.StringCodec;
import org.redisson.config.Config;

import java.security.SecureRandom;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * redis 启动方式：./src/redis-server redis.conf
 * @author 青石路
 * @date 2021/6/11 20:06
 */
public class LuaDemo {

    private RedissonClient redissonClient;

    @Before
    public void before() {
        Config config = new Config();
        config.useSingleServer()
                .setAddress("redis://192.168.1.110:6379");
        redissonClient = Redisson.create(config);
    }

    @After
    public void after() {
        redissonClient.shutdown();
    }

    /**
     * @author 青石路
     * 如果不存在，则进行 set，并设置过期时间
     * 如果存在，则返回过期时间
     */
    @Test
    public void simpleLua() {
        String key = "name";
        String value = "lee";
        int expireTime = 30;
        String luaScript = "if (redis.call('exists', KEYS[1]) == 0) then " +
                "redis.call('set', KEYS[1], ARGV[1]); " +
                "redis.call('expire', KEYS[1], ARGV[2]); " +
                "return nil; " +
                "end; " +
                "return redis.call('ttl', KEYS[1]); ";
        RScript script = redissonClient.getScript(new StringCodec());
        Long ttl  = script.eval(RScript.Mode.READ_WRITE, luaScript, RScript.ReturnType.INTEGER,
                Lists.newArrayList(key), value, expireTime);
        if (null == ttl) {
            System.out.println("key 设置成功, 过期时间：" + expireTime + " 秒");
        } else {
            System.out.println("key：" + key + " 已经存在, 过期时间还剩：" + ttl + " 秒");
        }
    }

    /**
     * 模拟 10 个线程获取锁：dist_lock
     */
    @Test
    public void distLockTest() {

        String id = UUID.randomUUID().toString();
        int lockExpireTime = 30000; // 锁有效时长，单位毫秒
        String lockName = "dist_lock";

        SecureRandom random = new SecureRandom();

        CountDownLatch latch = new CountDownLatch(10);
        for (int i=1; i<=10; i++) {
            new Thread(() -> {
                int nextInt = random.nextInt(10);
                System.out.println(Thread.currentThread().getName() + " 睡眠 " + nextInt + " 秒");

                try {
                    TimeUnit.SECONDS.sleep(nextInt);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }

                Long ttl = lock(lockName, lockExpireTime, id, Thread.currentThread().getId());
                if (ttl == null) {
                    System.out.println(Thread.currentThread().getName() + " 获取到锁, 执行业务...");
                    try {
                        // 业务执行时长
                        TimeUnit.SECONDS.sleep(2);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    unlock(lockName, lockExpireTime, id, Thread.currentThread().getId());
                    System.out.println(Thread.currentThread().getName() + " 业务执行完毕, 成功释放锁");
                } else {
                    System.out.println(Thread.currentThread().getName() + " 获取锁失败，锁过期时间还剩：" + ttl + " 毫秒");
                }
                latch.countDown();
            }, "t" + i).start();
        }

        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        System.out.println("主线程 end...");
    }

    /**
     * 加锁操作
     *      支持重入
     *
     * @param lockName
     * @param lockExpireTime
     * @param id
     * @param threadId
     * @return
     */
    private Long lock(String lockName, int lockExpireTime, String id, long threadId) {
        String lockScript = "if (redis.call('exists', KEYS[1]) == 0) then " +
                "redis.call('hincrby', KEYS[1], ARGV[2], 1); " +
                "redis.call('pexpire', KEYS[1], ARGV[1]); " +
                "return nil; " +
                "end; " +
                "if (redis.call('hexists', KEYS[1], ARGV[2]) == 1) then " +
                "redis.call('hincrby', KEYS[1], ARGV[2], 1); " +
                "redis.call('pexpire', KEYS[1], ARGV[1]); " +
                "return nil; " +
                "end; " +
                "return redis.call('pttl', KEYS[1]);";
        RScript script = redissonClient.getScript(new StringCodec());

        return script.eval(RScript.Mode.READ_WRITE, lockScript, RScript.ReturnType.INTEGER,
                Lists.newArrayList(lockName), lockExpireTime, id + ":" + threadId);
    }


    /**
     * 释放锁操作
     *      需要考虑重入次数
     * @param lockName
     * @param lockExpireTime
     * @param id
     * @param threadId
     * @return
     */
    private Long unlock(String lockName, int lockExpireTime, String id, long threadId) {
        String unlockScript = "if (redis.call('hexists', KEYS[1], ARGV[2]) == 0) then " +
                "return nil;" +
                "end; " +
                "local counter = redis.call('hincrby', KEYS[1], ARGV[2], -1); " +
                "if (counter > 0) then " +
                "redis.call('pexpire', KEYS[1], ARGV[1]); " +
                "return 0; " +
                "else " +
                "redis.call('del', KEYS[1]); " +
                "return 1; " +
                "end; " +
                "return nil;";
        RScript script = redissonClient.getScript(new StringCodec());

        return script.eval(RScript.Mode.READ_WRITE, unlockScript, RScript.ReturnType.INTEGER,
                Lists.newArrayList(lockName), lockExpireTime, id + ":" + threadId);
    }
}
